#include "trie.h"

int getLevenshteinDistance(const QString& a, const QString& b)
{
    int m = a.length();
    int n = b.length();
    QVector<QVector<int>> d(m + 1, QVector<int>(n + 1));
    for (int i = 0; i <= m; ++i) d[i][0] = i;
    for (int j = 0; j <= n; ++j) d[0][j] = j;
    for (int i = 1; i <= m; ++i)
    {
        for (int j = 1; j <= n; ++j)
        {
            int cost = (a[i - 1] == b[j - 1]) ? 0 : 1;
            d[i][j] = std::min({ d[i - 1][j] + 1, d[i][j - 1] + 1, d[i - 1][j - 1] + cost });
            if (i > 1 && j > 1 && a[i - 1] == b[j - 2] && a[i - 2] == b[j - 1]) d[i][j] = std::min(d[i][j], d[i - 2][j - 2] + cost);
        }
    }
    return d[m][n];
}

int getMaxCost(const QString& input)
{
    int length = input.length();
    if (length <= 3)
    {
        return 1;
    }
    else
    {
        return std::round(std::log2(length));
    }
}

TrieNode::TrieNode() :value(QChar()), isEndOfWord(false), files(QVector<FileInfo>()) {}
TrieNode::TrieNode(QChar Value) :value(Value), isEndOfWord(false), files(QVector<FileInfo>()) {}
TrieNode::TrieNode(QChar Value, bool Judge, FileInfo File) :value(Value), isEndOfWord(Judge), files(QVector<FileInfo>()) { files.push_back(File); }

Trie::Trie() { root = QSharedPointer<TrieNode>::create(); }

QSharedPointer<TrieNode> Trie::GetRoot() {return root;}

void Trie::Insert(const QString& str, const FileInfo& file)
{
    QSharedPointer<TrieNode> currentNode = root;
    for (int i = 0; i <= str.length() - 1; ++i)
    {
        QChar ch = str.at(i);
        if (!currentNode->children.contains(ch))
        {
            currentNode->children[ch] = QSharedPointer<TrieNode>::create(ch);
        }
        currentNode = currentNode->children[ch];
    }
    currentNode->isEndOfWord = true;
    currentNode->files.push_back(file);
}

void Trie::InsertWithSuffix(const QString& str, const FileInfo& file)
{
    if (str.isEmpty()) return;
    Insert(str, file);
    InsertWithSuffix(str.mid(1), file);
}

void Trie::collectWords(QSharedPointer<TrieNode> node, QVector<FileInfo>& result)
{
    if (node->isEndOfWord)
    {
        for (FileInfo& file : node->files)
            result.append(file);
    }
    for (QMap<QChar, QSharedPointer<TrieNode>>::iterator it = node->children.begin(); it != node->children.end(); ++it)
    {
        collectWords(it.value(), result);
    }
}

QVector<FileInfo> Trie::TrueSearchIgnoreCase(const QString& str, int index, QSharedPointer<TrieNode> node)
{
    if (index == str.length())
    {
        return node->isEndOfWord ? node->files : QVector<FileInfo>();
    }
    QVector<FileInfo> result;
    QChar lowerChar = str.at(index).toLower();
    QChar upperChar = str.at(index).toUpper();
    if (node->children.contains(lowerChar))
    {
        result += TrueSearchIgnoreCase(str, index + 1, node->children[lowerChar]);
    }
    if (lowerChar != upperChar && node->children.contains(upperChar))
    {
        result += TrueSearchIgnoreCase(str, index + 1, node->children[upperChar]);
    }
    return result;
}

QVector<FileInfo> Trie::SearchIgnoreCase(const QString& word)
{
    return TrueSearchIgnoreCase(word, 0, root);
}

QVector<FileInfo> Trie::TruePrefixSearchIgnoreCase(const QString& prefix, int index, QSharedPointer<TrieNode> node)
{
    if (index == prefix.length())
    {
        QVector<FileInfo> result;
        collectWords(node, result);
        return result;
    }
    QVector<FileInfo> result;
    QChar lowerChar = prefix.at(index).toLower();
    QChar upperChar = prefix.at(index).toUpper();
    if (node->children.contains(lowerChar))
    {
        result += TruePrefixSearchIgnoreCase(prefix, index + 1, node->children[lowerChar]);
    }
    if (lowerChar != upperChar && node->children.contains(upperChar))
    {
        result += TruePrefixSearchIgnoreCase(prefix, index + 1, node->children[upperChar]);
    }
    return result;
}

QVector<FileInfo> Trie::PrefixSearchIgnoreCase(const QString& prefix)
{
    return TruePrefixSearchIgnoreCase(prefix, 0, root);
}

QVector<FileInfo> Trie::Search(const QString& str)
{
    QSharedPointer<TrieNode> currentNode = root;
    for (int i = 0; i <= str.length()-1; ++i)
    {
        QChar ch = str.at(i);
        if (!currentNode->children.contains(ch))
        {
            return QVector<FileInfo>();
        }
        currentNode = currentNode->children[ch];
    }
    return currentNode->isEndOfWord ? currentNode->files : QVector<FileInfo>();
}

QVector<FileInfo> Trie::PrefixSearch(const QString& prefix)
{
    QSharedPointer<TrieNode> currentNode = root;
    QVector<FileInfo> result;
    for (int i = 0; i <= prefix.length()-1; ++i)
    {
        QChar ch = prefix.at(i);
        if (!currentNode->children.contains(ch))
        {
            return result;
        }
        currentNode = currentNode->children[ch];
    }
    collectWords(currentNode, result);
    return result;
}

void FuzzyCollectWords(QSharedPointer<TrieNode> node, const QString& currentStr, const QString& target, int maxCost, QVector<FileInfo>& result)
{
    if (node->isEndOfWord)
    {
        int distance = getLevenshteinDistance(currentStr, target);
        if (distance <= maxCost)
        {
            for (FileInfo& file : node->files)
                result.append(file);
        }
    }
    for (QMap<QChar, QSharedPointer<TrieNode>>::iterator it = node->children.begin(); it != node->children.end(); ++it)
    {
        QString nextWord = currentStr + it.key();
        FuzzyCollectWords(it.value(), nextWord, target, maxCost, result);
    }
}

QVector<FileInfo> Trie::FuzzySearch(const QString& word, int maxCost)
{
    QVector<FileInfo> result;
    FuzzyCollectWords(root, "", word, maxCost, result);
    return result;
}

QVector<FileInfo> Trie::FuzzyPrefixSearch(const QString& prefix, int maxCost)
{
    QVector<FileInfo> result;
    QSharedPointer<TrieNode> currentNode = root;
    for (int i = 0; i < prefix.length(); ++i)
    {
        QChar ch = prefix.at(i);
        if (!currentNode->children.contains(ch))
        {
            return result;
        }
        currentNode = currentNode->children[ch];
    }
    FuzzyCollectWords(currentNode, prefix, prefix, maxCost, result);
    return result;
}
